package cn.wujiangbo.recommend;

import org.ansj.domain.Result;
import org.ansj.domain.Term;
import org.ansj.splitWord.analysis.ToAnalysis;
import java.util.*;

/**
 * <p>TF-IDF工具类</p>
 * 解释：TF-IDF是一种从文本中提取关键词的常用算法，它基于单词在文本中的出现频率和在整个语料库中出现的频率，对单词进行权重计算。TF表示单词在当前文本中的频率，IDF表示单词在整个语料库中的逆文档频率，其计算公式如下：
 *      TF = (单词在文本中的出现次数) / (文本中单词总数)
 *      IDF = log(语料库中文档总数 / 含有该单词的文档数)
 *      TF-IDF = TF * IDF
 */
public class TfIdfTools {

    /**
     * 对文本内容进行分词，获取所有分词
     *
     * @param str 需要分词的文本内容
     */
    public static List<String> ansj(String str) {
        Set<String> Nature = new HashSet<>();
        //只取n和vn两种词性 即名词和一部分动词
        Nature.add("n");
        Nature.add("vn");
        // 分词结果的一个封装，主要是一个List<Term>的terms
        Result result = ToAnalysis.parse(str);

        // 拿到terms
        List<Term> terms = result.getTerms();

        List<String> list = new ArrayList<>();
        for (int i = 0; i < terms.size(); i++) {
            String word = terms.get(i).getName(); // 拿到词
            String natureStr = terms.get(i).getNatureStr(); // 拿到词性
            if (Nature.contains(natureStr)) // 判断set里有没有该词性，只取自己需要的词
            {
                list.add(word);
            }
        }
        return list;
    }

    /**
     * 获取每个词的TF-IDF值
     *
     * @param docs 需要计算的词汇集合
     * @return
     */
    public static Map<String, Double> tfIdf(List<String> docs) {
        Map<String, Integer> wordFreq = new HashMap<>();
        int totalWords = 0;
        for (String doc : docs) {
            String[] words = doc.split(" ");
            for (String word : words) {
                wordFreq.put(word, wordFreq.getOrDefault(word, 0) + 1);
                totalWords++;
            }
        }
        Map<String, Double> tfIdf = new HashMap<>();
        int docSize = docs.size();
        for (String word : wordFreq.keySet()) {
            double tf = (double) wordFreq.get(word) / totalWords;
            int docCount = 0;
            for (String doc : docs) {
                if (doc.contains(word)) {
                    docCount++;
                }
            }
            double idf = Math.log((double) docSize / (docCount + 1));
            tfIdf.put(word, tf * idf);
        }
        return tfIdf;
    }
}
